'''
Chaoyi 希望文章中的 Table 也能被收集.

https://www.ncbi.nlm.nih.gov/pmc/articles/PMC7372203/
https://www.ncbi.nlm.nih.gov/pmc/articles/PMC7804217/

1. 用爬虫把 Table 对应的 html 采集下来
2. 然后写 css 样式在本地进行渲染
3. 最后用 html -> png 的工具转化为图片

python src/fetch_table.py --extraction-dir /remote-home/share/medical/public/PMC_OA --volumes 4

'''
import glob
import os
import pathlib
import subprocess
import shutil
from tqdm import tqdm
import pandas as pd

import logging
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)
# 设置两个处理器handler
console_handler = logging.StreamHandler()
# 给两个相同名称的logger添加上处理器
logger.addHandler(console_handler)
# 设置一下格式
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - \33[32m%(message)s\033[0m')
console_handler.setFormatter(formatter)

from data import OA_LINKS  # file links of PMC Open Access
from args import parse_args_oa
from parser import get_volume_info, xml_to_table
from utils import read_jsonl, write_jsonl


def extract_tabs(volumes, extraction_dir: pathlib.Path):

    if not isinstance(extraction_dir, pathlib.Path):
        extraction_dir = pathlib.Path(extraction_dir)

    tabfig_records, tab_records = [], []

    for volume_id in volumes:
        volume = 'PMC00%dxxxxxx' % volume_id
        file_name='oa_comm_xml.%s.baseline.2022-09-03.filelist.csv' % volume
        file_path = extraction_dir / volume / file_name
        df = pd.read_csv(file_path, sep=',')

        for idx in tqdm(range(len(df)), desc='parse xml'):
            xml_path = extraction_dir / volume / df.loc[idx, 'Article File']
            tabfig_records += xml_to_table(xml_path)[0]
            tab_records += xml_to_table(xml_path)[1]
    # endfor
    # raise RuntimeError(f"tabfig: {len(tabfig_records)}; tab: {len(tab_records)}")

    return tabfig_records, tab_records


if __name__ == '__main__':
    # Check if wget is available
    if not shutil.which("wget"):
        print("wget not found, please install wget and put it on your PATH")
        exit(-1)
    args = parse_args_oa()

    # xml_to_table('/remote-home/weixionglin/build-pubmed/PMC554762.xml')

    save_name = ''.join([str(volume_id) for volume_id in args.volumes])

    tabs_path = f'{args.extraction_dir}/table/tab_{save_name}.jsonl'
    tabfigs_path = f'{args.extraction_dir}/table/tabfig_{save_name}.jsonl'

    if not os.path.exists(tabs_path):
        # Parse XML files into image info
        logger.info('Extracting Volume INFO')
        tabfigs_info, tabs_info = extract_tabs(
            volumes=args.volumes,
            extraction_dir=args.extraction_dir
        )

        # Save Volume info in jsonl
        logger.info('Saving Volume INFO')
        write_jsonl(data_list=tabs_info, save_path=tabs_path)  # tables
        write_jsonl(data_list=tabfigs_info, save_path=tabfigs_path)  # table figures

        logger.info('Saved')
    else:
        volume_info = read_jsonl(file_path=tabs_path)

    logger.info('Done')
